#!/usr/bin/env python3
"""
生成3D模型上传测试数据的脚本
用于为管理者提供可修改和保存的测试数据
"""

import asyncio
import sys
from datetime import datetime, timedelta
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from database.models import User, Product, Model3DUpload
from database.connection import get_session, init_db
from sqlmodel import select


async def create_test_users_and_products():
    """创建测试用户和产品数据"""
    async with get_session() as session:
        # 检查是否已存在测试用户
        result = await session.execute(select(User).where(User.username == "test_admin"))
        test_user = result.scalar_one_or_none()
        
        if not test_user:
            # 创建测试用户
            test_user = User(
                username="test_admin",
                email="test_admin@example.com",
                full_name="测试管理员",
                is_active=True,
                created_at=datetime.utcnow()
            )
            session.add(test_user)
            await session.commit()
            await session.refresh(test_user)
            print(f"创建测试用户: {test_user.full_name} (ID: {test_user.id})")
        else:
            print(f"使用现有测试用户: {test_user.full_name} (ID: {test_user.id})")
        
        # 检查是否已存在测试产品
        result = await session.execute(select(Product).where(Product.name.like("%测试产品%")))
        test_products = result.scalars().all()
        
        if len(test_products) == 0:
            # 创建测试产品
            test_product = Product(
                name="测试产品 - 3D模型上传",
                description="用于测试3D模型上传功能的产品",
                price=99.99,
                category="测试分类",
                is_active=True,
                created_at=datetime.utcnow()
            )
            session.add(test_product)
            await session.commit()
            await session.refresh(test_product)
            print(f"创建测试产品: {test_product.name} (ID: {test_product.id})")
        else:
            test_product = test_products[0]
            print(f"使用现有测试产品: {test_product.name} (ID: {test_product.id})")
        
        return test_user, test_product


async def create_test_model3d_uploads():
    """创建3D模型上传测试数据"""
    async with get_session() as session:
        # 获取或创建测试用户和产品
        test_user, test_product = await create_test_users_and_products()
        
        # 定义测试文件数据
        test_files = [
            {
                "original_filename": "test_cube_v1.stl",
                "file_url": "https://cdn.example.com/uploads/3dmodels/test_cube_v1.stl",
                "file_size": 1024 * 15,  # 15KB
                "file_format": "STL",
                "mime_type": "application/sla",
                "storage_path": "/uploads/3dmodels/2024/01/",
                "status": "UPLOADED",
                "generation_params": '{"quality": "high", "resolution": "medium"}',
                "ai_model_preference": "stable-diffusion-v1-5"
            },
            {
                "original_filename": "test_sphere_v2.obj",
                "file_url": "https://cdn.example.com/uploads/3dmodels/test_sphere_v2.obj",
                "file_size": 1024 * 32,  # 32KB
                "file_format": "OBJ",
                "mime_type": "text/plain",
                "storage_path": "/uploads/3dmodels/2024/01/",
                "status": "PROCESSING",
                "generation_params": '{"quality": "medium", "resolution": "low"}',
                "ai_model_preference": "midjourney-v4"
            },
            {
                "original_filename": "test_custom_design.stl",
                "file_url": "https://cdn.example.com/uploads/3dmodels/test_custom_design.stl",
                "file_size": 1024 * 128,  # 128KB
                "file_format": "STL",
                "mime_type": "application/sla",
                "storage_path": "/uploads/3dmodels/2024/01/",
                "status": "COMPLETED",
                "generation_params": '{"quality": "ultra", "resolution": "high", "texture": true}',
                "ai_model_preference": "dall-e-3",
                "processing_time": 45,
                "processing_started_at": datetime.utcnow() - timedelta(minutes=50),
                "processing_completed_at": datetime.utcnow() - timedelta(minutes=5)
            }
        ]
        
        created_count = 0
        
        for i, file_data in enumerate(test_files):
            # 检查是否已存在相同的测试文件
            result = await session.execute(
                select(Model3DUpload).where(
                    Model3DUpload.original_filename == file_data["original_filename"],
                    Model3DUpload.user_id == test_user.id
                )
            )
            existing_upload = result.scalar_one_or_none()
            
            if existing_upload:
                print(f"跳过已存在的上传记录: {file_data['original_filename']} (ID: {existing_upload.id})")
                continue
            
            # 创建新的3D模型上传记录
            upload = Model3DUpload(
                user_id=test_user.id,
                product_id=test_product.id,
                original_filename=file_data["original_filename"],
                file_url=file_data["file_url"],
                file_size=file_data["file_size"],
                file_format=file_data["file_format"],
                mime_type=file_data["mime_type"],
                storage_path=file_data["storage_path"],
                status=file_data["status"],
                generation_params=file_data.get("generation_params"),
                ai_model_preference=file_data.get("ai_model_preference"),
                processing_time=file_data.get("processing_time"),
                processing_started_at=file_data.get("processing_started_at"),
                processing_completed_at=file_data.get("processing_completed_at"),
                created_at=datetime.utcnow() - timedelta(days=i),  # 错开创建时间
                updated_at=datetime.utcnow() - timedelta(days=i)
            )
            
            session.add(upload)
            await session.commit()
            await session.refresh(upload)
            
            created_count += 1
            print(f"创建测试上传记录: {upload.original_filename} (ID: {upload.id})")
        
        print(f"\n测试数据生成完成！共创建 {created_count} 条新的3D模型上传记录")
        
        # 显示所有测试记录
        result = await session.execute(
            select(Model3DUpload).where(Model3DUpload.user_id == test_user.id)
        )
        all_uploads = result.scalars().all()
        
        print(f"\n测试用户 '{test_user.full_name}' 的所有3D模型上传记录:")
        print("-" * 80)
        for upload in all_uploads:
            print(f"ID: {upload.id:3} | 文件名: {upload.original_filename:25} | 状态: {upload.status:12} | 大小: {upload.file_size/1024:6.1f}KB | 创建时间: {upload.created_at.strftime('%Y-%m-%d %H:%M')}")
        
        print("\n" + "=" * 80)
        print("测试数据已准备就绪！管理员现在可以:")
        print("1. 在后台管理界面查看和编辑这些测试记录")
        print("2. 修改任何字段信息并保存")
        print("3. 测试各种状态和配置组合")
        print("4. 验证数据完整性和约束条件")
        print("=" * 80)


async def main():
    """主函数"""
    print("开始生成3D模型上传测试数据...")
    print("=" * 50)
    
    try:
        # 初始化数据库
        await init_db()
        print("数据库初始化完成")
        
        # 创建测试数据
        await create_test_model3d_uploads()
        
        print("\n✅ 测试数据生成成功！")
        
    except Exception as e:
        print(f"\n❌ 生成测试数据时出错: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(main())